/* Implements a Marquardt-Levenberg least-squares minimization to
   find a set of parameters p to fit a nonlinear function f(p, x)
   to a data set.

   [FP, CHISQ, NITER] = MRQFIT(F, P, X, Y) fits the function
   F(P, X) to the data specified by X and Y.  The parameters found
   by the minimization algorithm are returned to FP, the fit
   chi^2 normalized by the D.O.F. (data points - varying parameters)
   is returned to CHISQ, and the number of successful iterations
   is returned to NITER.
   
   [...] = MRQFIT(F, P, X, Y, SIG, VP, LB, UB, IMAX, TOL) provides
   additional information for controlling the minimization algorithm.
   SIG provides standard deviations for each data value in Y
   (the default is 1).
   VP indicates which parameters vary (default is all vary).
   LB and UB provide lower and upper bounds on the fit parameters
   (the default is none).  Use NaN or Inf to specify parameters
   without lower or upper bounds.
   IMAX is a scalar containing the maximum number of iterations
   (the default is 25).
   TOL is a scalar containing the chi^2 convergence tolerance
   (the default is 1.0e-3).
   Pass an empty array [] for any of these values to use the default
   values.

   [FP, CHISQ, NITER, FITC, ERR, DEP] = MRQFIT(...) returns the
   resulting fit FCN(P, X) evaluated at each X (FITC), the number
   of successful iterations (NITER), the estimated error on each
   fit parameter (ERR), and the fit parameter dependencies (DEP). */

/* Ignores imaginary parts of complex inputs */

/* All numeric inputs converted to DOUBLE if needed */

/* By:   S.C. Molitor (smolitor@bme.jhu.edu)
   Date: February 22, 1999 */


// inclusions

#include <math.h>
#include <string.h>
#include "function.h"
#include "mex.h"
#include "uintutil.h"


// default values

#define DEF_SIG		(1.0)			// data standard deviation
#define DEF_VARY	(1)				// parameter varies
#define	DEF_LB		(mxGetInf())	// no lower bound
#define	DEF_UB		(mxGetInf())	// no upper bound
#define	DEF_MAX		(25)			// maximum iterations
#define	DEF_TOL		(1.0e-3)		// chi^2 convergence tolerance


// function prototypes


// CREATE2DARRAY - allocate space for 2D array

double **Create2DArray(int	m,		// number of rows
					   int	n);		// number of columns

// DESTROY2DARRAY - deallocate space from 2D array

void Destroy2DArray(double	**array,	// pointer to array
					int		m);			// number of rows

// GETDOUBLEARRAY - obtain pointer to input data

double *GetDoubleArray(const mxArray	*input,		// input array
					   int				npts,		// number of input data elements
					   double			**prfree,	// array of pointers to free
					   int				*nfree);	// number of pointers to free

// INITDOUBLEARRAY - create pointer to input data if needed

double *InitDoubleArray(const mxArray	*prhs[],	// pointers to command line inputs
						int				nrhs,		// number of command line inputs
						int				irhs,		// number of desired input
						int				npts,		// length of desired input
						char			*errmsg,	// error message if improper length
						double			defval,		// value to initialize array with
						double			**prfree,	// array of pointers to free
						int				*nfree);	// number of pointers to free

// INITINTEGERARRAY - create pointer to input data if needed

int *InitIntegerArray(const mxArray	*prhs[],	// pointers to command line inputs
					  int			nrhs,		// number of command line inputs
					  int			irhs,		// number of desired input
					  int			npts,		// length of desired input
					  char			*errmsg,	// error message if improper length
					  int			defval,		// value to initialize array with
					  double		**prfree,	// array of pointers to free
					  int			*nfree);	// number of pointers to free


// MRQMINBD - perform Marquardt-Levenberg least-squares minimization
//            caller must provide f(p, x) in the form y = F(x, p[], dy/dp, np)

void mrqminbd(double	x[],		// x values of dataset
			  double	y[],		// y values of dataset
			  double	sig[],		// individual std. deviations
			  int		npts,		// number of data points
			  double	p[],		// function parameters p
			  double	lb[],		// lower bounds on function parameters
			  double	ub[],		// upper bounds on function parameters
			  int		npar,		// number of function parameters
			  int		ip[],		// indices of parameters that vary
			  int		nfit,		// number of parameters that vary
			  double	**covar,	// returned covariance matrix
			  double	**alpha,	// returned curvature matrix
			  double	beta[],		// returned chi^2 partials matrix
			  double	*chisq,		// returned chi^2 value
			  double	(*funcs)(double,	// function x value, returns f(p, x)
								 double [],	// function parameters p
								 double [],	// returned df/dp values
								 int),		// number of parameters
			  double	*lambda);	// returned curvature diagonal scaling factor


// MEXFUNCTION - gateway routine

void mexFunction(int			nlhs,
				 mxArray		*plhs[],
				 int			nrhs,
				 const mxArray	*prhs[]) {

	// declare variables

	char		*sfcn;				// name of function F(p, x)
	double		**prfree;			// array of pointers to free at end
	double		*p, *x, *y, *sig;	// parameters p, data x & y, std. dev. of y
	double		*lb, *ub, *tol;		// lower & upper bounds, chi^2 tolerance
	double		**alpha, *beta;		// curvature matrix & chi^2 partials
	double		**covar;			// covariance matrix
	double		*fp;				// fit parameters for output
	double		lambda, olambda;	// covariance diagonal scaling
	double		*niter, *fitc;		// number of iterations & F(p, x)
	double		*chisq, ochisq;		// chi^2 values from fit
	double		*perr, *pdep;		// parameter errors & dependencies
	int			nfree;				// number of pointers to free at end
	int			nfcn;				// number of characters in function name
	int			npar, npts;			// number of parameters & data points
	int			*vp, *imax;			// flags to vary parameters & max iterations
	int			nfit;				// number of parameters that vary
	int			*ip;				// indices of parameters that vary
	int			i;					// loop counter

	// declare function F(p, x)

	double		(*fcn)(double,		// function x value, returns f(p, x)
					   double [],	// function parameters p
					   double [],	// returned df/dp values
					   int);		// number of parameters

	// allocate space to free arrays if needed

	prfree = (double **) mxMalloc((size_t) nrhs * sizeof(double *));
	nfree = 0;

	// call VALIDATEINPUTS to validate calling syntax
	// obtain pointer to input Y

	ValidateInputs(nlhs, nrhs, prhs, 4, 10, 6);

	// obtain function name

	nfcn = mxGetNumberOfElements(prhs[0]);
	if (nfcn < 1) {
		mexErrMsgTxt("FCN cannot be empty");
	}
	else if (mxGetClassID(prhs[0]) != mxCHAR_CLASS) {
		mexErrMsgTxt("FCN must be a string");
	}
	else {
		nfcn++;
		sfcn = (char *) mxCalloc(nfcn, sizeof(char));
		nfcn = mxGetString(prhs[0], sfcn, nfcn);
	}

	// set pointer to function if valid name

	if (strcmp(sfcn, "boltzmann") == 0) {
		fcn = oneboltzmann;
	}
	else if (strcmp(sfcn, "sum_boltzmann") == 0) {
		fcn = sumboltzmann;
	}
	else if (strcmp(sfcn, "curr_boltzmann") == 0) {
		fcn = currboltzmann;
	}
	else if (strcmp(sfcn, "exponential") == 0) {
		fcn = exponential;
	}
	else if (strcmp(sfcn, "gaussian") == 0) {
		fcn = gaussian;
	}
	else {
		mexErrMsgTxt("Invalid FCN name");
	}
	mxFree((void *) sfcn);

	// obtain number of parameters & data points

	npar = mxGetNumberOfElements(prhs[1]);
	if (npar < 1) {
		mexErrMsgTxt("P cannot be empty");
	}
	npts = mxGetNumberOfElements(prhs[2]);
	if (npts != mxGetNumberOfElements(prhs[3])) {
		mexErrMsgTxt("X and Y must have the same length");
	}
	else if (npts < 1) {
		mexErrMsgTxt("X and Y cannot be empty");
	}

	// obtain pointers to function parameters & data
	// convert these to DOUBLE if needed

	p = GetDoubleArray(prhs[1], npar, prfree, &nfree);
	x = GetDoubleArray(prhs[2], npts, prfree, &nfree);
	y = GetDoubleArray(prhs[3], npts, prfree, &nfree);

	// obtain other command line inputs if they exist
	// otherwise set to default value

	sig = InitDoubleArray(prhs, nrhs, 4, npts, "SIG must have same length as X and Y", DEF_SIG, prfree, &nfree);
	vp = InitIntegerArray(prhs, nrhs, 5, npar, "VP must have same length as P", DEF_VARY, prfree, &nfree);
	lb = InitDoubleArray(prhs, nrhs, 6, npar, "LB must have same length as P", DEF_LB, prfree, &nfree);
	ub = InitDoubleArray(prhs, nrhs, 7, npar, "UB must have same length as P", DEF_UB, prfree, &nfree);
	imax = InitIntegerArray(prhs, nrhs, 8, 1, "IMAX must be scalar", DEF_MAX, prfree, &nfree);
	tol = InitDoubleArray(prhs, nrhs, 9, 1, "TOL must be scalar", DEF_TOL, prfree, &nfree);

	// allocate arrays for curve fitting procedure

	ip = (int *) mxCalloc(npar, sizeof(int));
	alpha = Create2DArray(npar, npar);
	beta = (double *) mxCalloc(npar, sizeof(double));
	covar = Create2DArray(npar, npar);

	// create outputs for FP, CHISQ & NITER

	plhs[0] = mxCreateDoubleMatrix(1, npar, mxREAL);
	fp = mxGetPr(plhs[0]);
	plhs[1] = mxCreateDoubleMatrix(1, 1, mxREAL);
	chisq = mxGetPr(plhs[1]);
	plhs[2] = mxCreateDoubleMatrix(1, 1, mxREAL);
	niter = mxGetPr(plhs[2]);

	// copy input P to output FP
	// determine which parameters vary

	for (nfit = 0, i = 0; i < npar; i++) {
		fp[i] = p[i];
		if (vp[i] > 0) {
			ip[nfit] = i;
			nfit++;
		}
	}

	// call mrqminbd.c with negative lambda for initialization
	// save chisq & lambda values
	
	lambda = -1.0;
	*chisq = 0.0;
	mrqminbd(x, y, sig, npts, fp, lb, ub, npar, ip, nfit, covar, alpha, beta, chisq, fcn, &lambda);
	ochisq = *chisq;
	olambda = lambda;

	// make sure initial step results in a decrement of lambda
	// if not, keep trying until appropriate lambda is found

	for (i = 0; i < *imax; i++) {
		mrqminbd(x, y, sig, npts, fp, lb, ub, npar, ip, nfit, covar, alpha, beta, chisq, fcn, &lambda);
		if (lambda < olambda) {
			break;
		}
		else {
			olambda = lambda;
		}
	}

	// if appropriate starting lambda was found, iterate until convergence
	// make sure lambda is decreasing at convergence, not increasing!

	if (i < *imax) {
		for (i = 0; i < *imax; i++) {
			if (*chisq <= 0.0) {
				break;
			}
			if ((lambda < olambda) && (fabs((*chisq - ochisq) / *chisq) < *tol)) {
				break;
			}
			else {
				ochisq = *chisq;
				olambda = lambda;
			}
			mrqminbd(x, y, sig, npts, fp, lb, ub, npar, ip, nfit, covar, alpha, beta, chisq, fcn, &lambda);
		}
	}

	// otherwise, provide a warning message
	// original values will be returned

	else {
		mexWarnMsgTxt("Could not reduce chi^2 within maximum iterations");
	}

	// obtain results & free allocated space in mrqminbd.c

	lambda = 0.0;
	mrqminbd(x, y, sig, npts, fp, lb, ub, npar, ip, nfit, covar, alpha, beta, chisq, fcn, &lambda);
	*niter = (double) i;

	// calculate fit curve if specified

	if (nlhs > 3) {
		plhs[3] = mxCreateDoubleMatrix(1, npts, mxREAL);
		fitc = mxGetPr(plhs[3]);
		for (i = 0; i < npts; i++) {
			fitc[i] = (*fcn)(x[i], fp, beta, npar);
		}
	}

	// calculate parameter errors if specified

	if (nlhs > 4) {
		plhs[4] = mxCreateDoubleMatrix(1, npar, mxREAL);
		perr = mxGetPr(plhs[4]);
		for (i = 0; i < npar; i++) {
			if (covar[i][i] >= 0.0) {
				perr[i] = sqrt(covar[i][i] * (*chisq));
			}
			else {
				perr[i] = mxGetNaN();
			}
		}
	}

	// calculate parameter dependencies if specified

	if (nlhs > 5) {
		plhs[5] = mxCreateDoubleMatrix(1, npar, mxREAL);
		pdep = mxGetPr(plhs[5]);
		for (i = 0; i < npar; i++) {
			if ((alpha[i][i] > 0.0) && (covar[i][i] > 0.0)) {
				pdep[i] = 1.0 - (1.0 / (alpha[i][i] * covar[i][i]));
			}
			else {
				pdep[i] = mxGetNaN();
			}
		}
	}

	// destroy created arrays
	// free allocated inputs

	Destroy2DArray(covar, npar);
	mxFree((void *) beta);
	Destroy2DArray(alpha, npar);
	mxFree((void *) ip);
	for (i = nfree - 1; i >= 0; i--) {
		mxFree((void *) prfree[i]);
	}
	mxFree((void *) prfree);
}


// CREATE2DARRAY - allocate space for 2D array

double **Create2DArray(int	m,		// number of rows
					   int	n) {	// number of columns

	// declare variables

	double	**array;		// pointer to array
	int		i;				// loop counter

	// allocate space for 2D array

	array = (double **) mxMalloc((size_t) m * sizeof(double *));
	for (i = 0; i < m; i++) {
		array[i] = (double *) mxCalloc(n, sizeof(double));
	}
	return(array);
}


// DESTROY2DARRAY - deallocate space from 2D array

void Destroy2DArray(double	**array,	// pointer to array
					int		m) {		// number of rows

	// declare variables

	int		i;		// loop counter

	// deallocate space from 2D array

	for (i = 0; i < m; i++) {
		mxFree((void *) array[i]);
	}
	mxFree((void *) array);
}


// GETDOUBLEARRAY - obtain pointer to input data

double *GetDoubleArray(const mxArray	*input,		// input array
					   int				npts,		// number of input data elements
					   double			**prfree,	// array of pointers to free
					   int				*nfree) {	// number of pointers to free

	// declare variables

	mxClassID	classID;	// input class ID
	double		*prinput;	// pointer to data in input
	int			i;			// loop counter

	// if already DOUBLE, obtain pointer

	classID = mxGetClassID(input);
	if (classID == mxDOUBLE_CLASS) {
		prinput = mxGetPr(input);
	}

	// otherwise convert to DOUBLE
	// add pointer to list for mxFree later on

	else {
		prinput = (double *) mxCalloc(npts, sizeof(double));
		for (i = 0; i < npts; i++) {
			prinput[i] = GetDouble(input, classID, i);
		}
		prfree[*nfree] = prinput;
		(*nfree)++;
	}
	return(prinput);
}


// INITDOUBLEARRAY - create pointer to input data if needed

double *InitDoubleArray(const mxArray	*prhs[],	// pointers to command line inputs
						int				nrhs,		// number of command line inputs
						int				irhs,		// number of desired input
						int				npts,		// length of desired input
						char			*errmsg,	// error message if improper length
						double			defval,		// value to initialize array with
						double			**prfree,	// array of pointers to free
						int				*nfree) {	// number of pointers to free

	// declare variables

	bool	initflag;	// flag for initializing array
	double	*prinput;	// pointer to data in input
	int		i;			// loop counter

	// get pointer if input is provided & has correct number of points
	// empty input will be initialized to default value
	// display error message only if input is not empty

	if (nrhs > irhs) {
		if (mxIsEmpty(prhs[irhs])) {
			initflag = true;
		}
		else if (npts == mxGetNumberOfElements(prhs[irhs])) {
			prinput = GetDoubleArray(prhs[irhs], npts, prfree, nfree);
			initflag = false;
		}
		else {
			mexErrMsgTxt(errmsg);
		}
	}

	// if input not provided, will be initialized to default value

	else {
		initflag = true;
	}

	// create DOUBLE array if needed
	// initialize values to default value

	if (initflag) {
		prinput = (double *) mxCalloc(npts, sizeof(double));
		for (i = 0; i < npts; i++) {
			prinput[i] = defval;
		}
		prfree[*nfree] = prinput;
		(*nfree)++;
	}
	return(prinput);
}


// INITINTEGERARRAY - create pointer to input data if needed

int *InitIntegerArray(const mxArray	*prhs[],	// pointers to command line inputs
					  int			nrhs,		// number of command line inputs
					  int			irhs,		// number of desired input
					  int			npts,		// length of desired input
					  char			*errmsg,	// error message if improper length
					  int			defval,		// value to initialize array with
					  double		**prfree,	// array of pointers to free
					  int			*nfree) {	// number of pointers to free

	// declare variables

	bool		initflag;	// flag for initializing array
	int			*prinput;	// pointer to data in input
	int			i;			// loop counter
	mxClassID	classID;	// input class ID

	// input with correct number of points will be loaded as INT
	// empty input will be initialized to default value
	// display error message only if input is not empty

	if (nrhs > irhs) {
		if (mxIsEmpty(prhs[irhs])) {
			initflag = true;
		}
		else if (npts == mxGetNumberOfElements(prhs[irhs])) {
			classID = mxGetClassID(prhs[irhs]);
			initflag = false;
		}
		else {
			mexErrMsgTxt(errmsg);
		}
	}

	// if input not provided, will be initialized to default value

	else {
		initflag = true;
	}

	// create INT array
	// initialize values to default value if specified
	// otherwise, load integer values

	prinput = (int *) mxCalloc(npts, sizeof(int));
	for (i = 0; i < npts; i++) {
		if (initflag) {
			prinput[i] = defval;
		}
		else {
			prinput[i] = (int) GetDouble(prhs[irhs], classID, i);
		}
	}
	
	// add array to list to free at end
	// return pointer to array

	prfree[*nfree] = (double *) prinput;
	(*nfree)++;
	return(prinput);
}



